from pathlib import Path
import numpy as np
from tqdm import tqdm

from utils import process_record, save_record


DATA_ROOT = Path('datasets/mit-bih-supraventricular-arrhythmia-database-1.0.0')
OUTPUT_FOLDER = Path('processed_data/mit-bih-supra/')
LABEL_MAPPING = {
    "N": 0, "L": 0, "R": 0, "e": 0, "j": 0,  # N
    "A": 1, "a": 1, "S": 1, "J": 1,  # SVEB
    "V": 2, "E": 2,  # VEB
    "F": 3,  # F
    "/": 4, "f": 4, "Q": 4  # Q
}


if __name__ == '__main__':
    patient_ids = [int(p.stem) for p in DATA_ROOT.glob('*.hea')]
    for patient_id in tqdm(patient_ids):
        result = process_record(DATA_ROOT, patient_id, LABEL_MAPPING, lead='ECG2')
        if result is not None:
            patient_folder = OUTPUT_FOLDER / f"PATIENT_{patient_id}" / 'RECORD_0'
            save_record(patient_folder, *result)
            

    
    
